//------------------------------------------------------------------------------
// Categorizers.cpp
//------------------------------------------------------------------------------
// File provided for Microsoft Shared Source.
// Copyright (c) Microsoft Corporation. All rights reserved.
//
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF 
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO 
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A 
// PARTICULAR PURPOSE.
//------------------------------------------------------------------------------
//
// Helper class to assemble result sets into categories
//
//------------------------------------------------------------------------------

#include "pch.h"
#include "SortableArray.h"

//------------------------------------------------------------------------------
// CUnknownArray
//------------------------------------------------------------------------------

class ATL_NO_VTABLE CUnknownArray :
    public CComObjectRootEx<CComSingleThreadModel>,
    public IUnknownArray
{
public:
    static HRESULT s_CreateInstance(const CAtlArray<CComPtr<IUnknown>> &rgItems, REFIID riid, __deref_out void **ppv);

    BEGIN_COM_MAP(CUnknownArray)
        COM_INTERFACE_ENTRY(IUnknownArray)
    END_COM_MAP()

    DECLARE_PROTECT_FINAL_CONSTRUCT()

    // IUnknownArray
    IFACEMETHODIMP GetCount(__out DWORD *pcItems);
    IFACEMETHODIMP GetItem(DWORD dwIndex, REFIID riid, __deref_out void **ppv);

private:
    HRESULT _Initialize(const CAtlArray<CComPtr<IUnknown>> &rgItems);

private:
    CAtlArray<CComPtr<IUnknown>> _rgItems;
};

HRESULT CUnknownArray::_Initialize(const CAtlArray<CComPtr<IUnknown>> &rgItems)
{
    HRESULT hr = S_OK;
    try
    {
        _rgItems.Copy(rgItems);
    }
    catch (CAtlException &e)
    {
        hr = e.m_hr;
    }
    return hr;
}

HRESULT CUnknownArray::s_CreateInstance(const CAtlArray<CComPtr<IUnknown>> &rgItems, REFIID riid, __deref_out void **ppv)
{
    *ppv = NULL;
    CComObjectWithRef<CUnknownArray> *pspUnkArray;
    HRESULT hr = CComObjectWithRef<CUnknownArray>::CreateInstance(&pspUnkArray);
    if (SUCCEEDED(hr))
    {
        hr = pspUnkArray->_Initialize(rgItems);
        if (SUCCEEDED(hr))
        {
            hr = pspUnkArray->QueryInterface(riid, ppv);
        }
        pspUnkArray->Release();
    }
    return hr;
}

// IUnknownArray
STDMETHODIMP CUnknownArray::GetCount(__out DWORD *pcItems)
{
    *pcItems = (DWORD)_rgItems.GetCount();
    return S_OK;
}

STDMETHODIMP CUnknownArray::GetItem(DWORD dwIndex, REFIID riid, __deref_out void **ppv)
{
    HRESULT hr = E_INVALIDARG;
    if (dwIndex >= 0 && dwIndex < _rgItems.GetCount())
    {
        hr = _rgItems.GetAt(dwIndex)->QueryInterface(riid, ppv);
    }
    return hr;
}

//------------------------------------------------------------------------------
// Special traits class for doing custom comparsions of ISolutionItem
//------------------------------------------------------------------------------

struct SOLUTIONITEMSORTINFO
{
    COLUMNID colid;
    BOOL fAscending;
};

class CCompareSolutionItemTraits : public CElementTraits<CComPtr<ISolutionItem>>
{
public:
    static bool CompareElements(const CComPtr<ISolutionItem> &spsi1, const CComPtr<ISolutionItem> &spsi2)
    {
        return CompareElementsOrdered(spsi1, spsi2) == 0;
    }

    static int CompareElementsOrderedParam(const CComPtr<ISolutionItem> &spsi1, const CComPtr<ISolutionItem> &spsi2, LPARAM lParam)
    {
        const SOLUTIONITEMSORTINFO *psisi = (const SOLUTIONITEMSORTINFO *)lParam;
        CComPtr<ISolutionItem> spsiCmp1;
        CComPtr<ISolutionItem> spsiCmp2;

        spsiCmp1 = psisi->fAscending ? spsi1 : spsi2;
        spsiCmp2 = psisi->fAscending ? spsi2 : spsi1;

        int iCmp = 0;
        switch (psisi->colid)
        {
        case COLID_NAME:
            break;

        case COLID_PATH:
            iCmp = ComparePathOrdered(spsiCmp1, spsiCmp2);
            if (iCmp == 0)
            {
                iCmp = CompareNameOrdered(spsiCmp1, spsiCmp2);
            }
            break;

        case COLID_SIZE:
            iCmp = CompareSizeOrdered(spsiCmp1, spsiCmp2);
            break;

        case COLID_MODIFIEDDATE:
            iCmp = CompareModifiedOrdered(spsiCmp1, spsiCmp2);
            break;

        default:
            ATLASSERT(FALSE);
            break;
        }

        if (psisi->colid != COLID_PATH && iCmp == 0)
        {
            iCmp = CompareNameOrdered(spsiCmp1, spsiCmp2);
            if (iCmp == 0)
            {
                iCmp = ComparePathOrdered(spsiCmp1, spsiCmp2);
            }
        }
        return iCmp;
    }

    static int CompareNameOrdered(const CComPtr<ISolutionItem> &spsi1, const CComPtr<ISolutionItem> &spsi2)
    {
        int iRet = 0;
        WCHAR sz1[MAX_PATH];
        if (SUCCEEDED(spsi1->GetName(sz1, ARRAYSIZE(sz1))))
        {
            WCHAR sz2[MAX_PATH];
            if (SUCCEEDED(spsi2->GetName(sz2, ARRAYSIZE(sz2))))
            {
                iRet = _wcsicmp(sz1, sz2);
            }
        }
        return iRet;
    }

    static int ComparePathOrdered(const CComPtr<ISolutionItem> &spsi1, const CComPtr<ISolutionItem> &spsi2)
    {
        int iRet = 0;
        WCHAR sz1[MAX_PATH];
        if (SUCCEEDED(spsi1->GetPath(sz1, ARRAYSIZE(sz1))))
        {
            WCHAR sz2[MAX_PATH];
            if (SUCCEEDED(spsi2->GetPath(sz2, ARRAYSIZE(sz2))))
            {
                iRet = _wcsicmp(sz1, sz2);
            }
        }
        return iRet;
    }

    static int CompareSizeOrdered(const CComPtr<ISolutionItem> &spsi1, const CComPtr<ISolutionItem> &spsi2)
    {
        int iRet = 0;
        LARGE_INTEGER li1;
        if (SUCCEEDED(spsi1->GetSize(&li1)))
        {
            LARGE_INTEGER li2;
            if (SUCCEEDED(spsi2->GetSize(&li2)))
            {
                LONGLONG llRet = li1.QuadPart - li2.QuadPart;
                if (llRet == 0)
                {
                    iRet = 0;
                }
                else if (llRet < 0)
                {
                    iRet = -1;
                }
                else
                {
                    iRet = 1;
                }
            }
        }
        return iRet;
    }

    static int CompareModifiedOrdered(const CComPtr<ISolutionItem> &spsi1, const CComPtr<ISolutionItem> &spsi2)
    {
        int iRet = 0;
        FILETIME ft1;
        if (SUCCEEDED(spsi1->GetModified(&ft1)))
        {
            FILETIME ft2;
            if (SUCCEEDED(spsi2->GetModified(&ft2)))
            {
                iRet = CompareFileTime(&ft1, &ft2);
            }
        }
        return iRet;
    }
};

//------------------------------------------------------------------------------
// CNullCategorizer
//
// Categorizes items with no categories
//------------------------------------------------------------------------------

typedef CSortableArray<CComPtr<ISolutionItem>, CCompareSolutionItemTraits> CSolutionItemArray;

class ATL_NO_VTABLE CNullCategorizer : 
    public CComObjectRootEx<CComSingleThreadModel>,
    public CComCoClass<CNullCategorizer, &__uuidof(NullCategorizer)>,
    public ISolutionItemCategorizer
{
public:
    CNullCategorizer() {}

    DECLARE_NO_REGISTRY()
    DECLARE_NOT_AGGREGATABLE(CNullCategorizer)

    BEGIN_COM_MAP(CNullCategorizer)
        COM_INTERFACE_ENTRY(ISolutionItemCategorizer)
    END_COM_MAP()

    DECLARE_PROTECT_FINAL_CONSTRUCT()

    // ISolutionItemCategorizer
    IFACEMETHODIMP Initialize(COLUMNID colidSort, BOOL fSortAscending);
    IFACEMETHODIMP Categorize(INDEXENTRYTYPE iet, __in ISolutionItem *psi);
    IFACEMETHODIMP GetCategories(REFIID riid, __deref_out void **ppv);

private:
    COLUMNID _colidSort;
    BOOL _fSortAscending;
    CSolutionItemArray _rgItems;
};

OBJECT_ENTRY_AUTO(__uuidof(NullCategorizer), CNullCategorizer)

// ISolutionItemCategorizer
STDMETHODIMP CNullCategorizer::Initialize(COLUMNID colidSort, BOOL fSortAscending)
{
    _colidSort = colidSort;
    _fSortAscending = fSortAscending;
    return S_OK;
}

STDMETHODIMP CNullCategorizer::Categorize(INDEXENTRYTYPE iet, __in ISolutionItem *psi)
{
    HRESULT hr = S_OK;
    BOOL fFound;
    SOLUTIONITEMSORTINFO sisi = { _colidSort, _fSortAscending };
    size_t iInsert = _rgItems.Search(psi, TRUE, fFound, (LPARAM)&sisi);
    if (!fFound)
    {
        try
        {
            _rgItems.InsertAt(iInsert, psi);
        }
        catch (CAtlException &e)
        {
            hr = e.m_hr;
        }
    }
    return hr;
}

STDMETHODIMP CNullCategorizer::GetCategories(REFIID riid, __deref_out void **ppv)
{
    return CUnknownArray::s_CreateInstance((const CAtlArray<CComPtr<IUnknown>> &)_rgItems, riid, ppv);
}

//------------------------------------------------------------------------------
// CHitTypeCategorizer
//
// Categorizes solution items based on the INDEXENTRYTYPE
//------------------------------------------------------------------------------

class ATL_NO_VTABLE CHitTypeCategorizer : 
    public CComObjectRootEx<CComSingleThreadModel>,
    public CComCoClass<CHitTypeCategorizer, &__uuidof(HitTypeCategorizer)>,
    public ISolutionItemCategorizer
{
public:
    CHitTypeCategorizer() {}

    DECLARE_NO_REGISTRY()
    DECLARE_NOT_AGGREGATABLE(CHitTypeCategorizer)

    BEGIN_COM_MAP(CHitTypeCategorizer)
        COM_INTERFACE_ENTRY(ISolutionItemCategorizer)
    END_COM_MAP()

    DECLARE_PROTECT_FINAL_CONSTRUCT()

    // ISolutionItemCategorizer
    IFACEMETHODIMP Initialize(COLUMNID colidSort, BOOL fSortAscending);
    IFACEMETHODIMP Categorize(INDEXENTRYTYPE iet, __in ISolutionItem *psi);
    IFACEMETHODIMP GetCategories(REFIID riid, __deref_out void **ppv);

private:
    HRESULT _AddItemArrayGroupToGroupArray(UINT uiResIDName, const CSolutionItemArray &rgItems, __inout CAtlArray<CComPtr<ISolutionItemGroup>> &rgGroups);

private:
    COLUMNID _colidSort;
    BOOL _fSortAscending;
    CSolutionItemArray _rgFileItems;
    CSolutionItemArray _rgPathItems;
};

OBJECT_ENTRY_AUTO(__uuidof(HitTypeCategorizer), CHitTypeCategorizer)

// ISolutionItemCategorizer
STDMETHODIMP CHitTypeCategorizer::Initialize(COLUMNID colidSort, BOOL fSortAscending)
{
    _colidSort = colidSort;
    _fSortAscending = fSortAscending;
    return S_OK;
}

STDMETHODIMP CHitTypeCategorizer::Categorize(INDEXENTRYTYPE iet, __in ISolutionItem *psi)
{
    HRESULT hr = S_OK;
    try
    {
        BOOL fFound;
        size_t iInsert;
        SOLUTIONITEMSORTINFO sisi = { _colidSort, _fSortAscending };
        if (iet == IET_FILENAME)
        {
            // If this item is in the path part array, remove it.  File hits take precedence
            size_t i = _rgPathItems.Search(psi, FALSE, fFound, (LPARAM)&sisi);
            if (fFound)
            {
                _rgPathItems.RemoveAt(i);
            }

            iInsert = _rgFileItems.Search(psi, TRUE, fFound, (LPARAM)&sisi);
            if (!fFound)
            {
                _rgFileItems.InsertAt(iInsert, psi);
            }
        }
        else
        {
            // If this item is in the file array, ignore it.  File hits take precedence.
            // If this item is in the path part array, remove it.  File hits take precedence
            size_t i = _rgFileItems.Search(psi, FALSE, fFound, (LPARAM)&sisi);
            if (!fFound)
            {
                iInsert = _rgPathItems.Search(psi, TRUE, fFound, (LPARAM)&sisi);
                if (!fFound)
                {
                    _rgPathItems.InsertAt(iInsert, psi);
                }
            }
        }
    }
    catch (CAtlException &e)
    {
        hr = e.m_hr;
    }
    return hr;
}

HRESULT CHitTypeCategorizer::_AddItemArrayGroupToGroupArray(UINT uiResIDName, const CSolutionItemArray &rgItems, __inout CAtlArray<CComPtr<ISolutionItemGroup>> &rgGroups)
{
    CComPtr<IUnknownArray> spua;
    HRESULT hr = CUnknownArray::s_CreateInstance((const CAtlArray<CComPtr<IUnknown>> &)rgItems, IID_PPV_ARGS(&spua));
    if (SUCCEEDED(hr))
    {
        CString strName;
        hr = strName.LoadString(uiResIDName) ? S_OK : E_INVALIDARG;
        if (SUCCEEDED(hr))
        {
            CComPtr<ISolutionItemGroup> spsig;
            hr = CreateSolutionItemGroup(strName, spua, IID_PPV_ARGS(&spsig));
            if (SUCCEEDED(hr))
            {
                rgGroups.Add(spsig);
            }
        }
    }
    return hr;
}

STDMETHODIMP CHitTypeCategorizer::GetCategories(REFIID riid, __deref_out void **ppv)
{
    CAtlArray<CComPtr<ISolutionItemGroup>> rgGroups;
    HRESULT hr = S_OK;
    try
    {
        if (_rgFileItems.GetCount())
        {
            hr = _AddItemArrayGroupToGroupArray(IDS_HITTYPEFILE, _rgFileItems, rgGroups);
        }

        if (SUCCEEDED(hr) && _rgPathItems.GetCount())
        {
            hr = _AddItemArrayGroupToGroupArray(IDS_HITTYPEPATH, _rgPathItems, rgGroups);
        }
    }
    catch (CAtlException &e)
    {
        hr = e.m_hr;
    }

    if (SUCCEEDED(hr))
    {
        hr = CUnknownArray::s_CreateInstance((const CAtlArray<CComPtr<IUnknown>> &)rgGroups, riid, ppv);
    }

    return hr;
}

//------------------------------------------------------------------------------
// CTypeCategorizer
//
// Categorizes solution items based on file type
//------------------------------------------------------------------------------

class ATL_NO_VTABLE CTypeCategorizer : 
    public CComObjectRootEx<CComSingleThreadModel>,
    public CComCoClass<CNullCategorizer, &__uuidof(NullCategorizer)>,
    public ISolutionItemCategorizer
{
public:
    CTypeCategorizer() {}

    DECLARE_NO_REGISTRY()
    DECLARE_NOT_AGGREGATABLE(CTypeCategorizer)

    BEGIN_COM_MAP(CTypeCategorizer)
        COM_INTERFACE_ENTRY(ISolutionItemCategorizer)
    END_COM_MAP()

    DECLARE_PROTECT_FINAL_CONSTRUCT()

    // ISolutionItemCategorizer
    IFACEMETHODIMP Initialize(COLUMNID colidSort, BOOL fSortAscending);
    IFACEMETHODIMP Categorize(INDEXENTRYTYPE iet, __in ISolutionItem *psi);
    IFACEMETHODIMP GetCategories(REFIID riid, void **ppv);

private:
    HRESULT _GetTypeArrayForItem(__in ISolutionItem *psi, __out CSolutionItemArray **ppsia);
    HRESULT _CreateGroupForTypeArray(PCWSTR pszType, const CSolutionItemArray &sia, REFIID riid, __deref_out void **ppv);

private:
    COLUMNID _colidSort;
    BOOL _fSortAscending;
    CAtlMap<CString, CSolutionItemArray> _mapTypes;
};

OBJECT_ENTRY_AUTO(__uuidof(TypeCategorizer), CTypeCategorizer)

// ISolutionItemCategorizer
STDMETHODIMP CTypeCategorizer::Initialize(COLUMNID colidSort, BOOL fSortAscending)
{
    _colidSort = colidSort;
    _fSortAscending = fSortAscending;
    return S_OK;
}

HRESULT CTypeCategorizer::_GetTypeArrayForItem(__in ISolutionItem *psi, __out CSolutionItemArray **ppsia)
{
    WCHAR szType[MAX_PATH];
    HRESULT hr = psi->GetCanonicalType(szType, ARRAYSIZE(szType));
    if (SUCCEEDED(hr))
    {
        CSolutionItemArray &sia = _mapTypes[szType];
        *ppsia = &sia;
    }
    return hr;
}

STDMETHODIMP CTypeCategorizer::Categorize(INDEXENTRYTYPE iet, __in ISolutionItem *psi)
{
    HRESULT hr = E_FAIL;
    try
    {
        CSolutionItemArray *psia;
        hr = _GetTypeArrayForItem(psi, &psia);
        if (SUCCEEDED(hr))
        {
            BOOL fFound;
            SOLUTIONITEMSORTINFO sisi = { _colidSort, _fSortAscending };
            size_t iInsert = psia->Search(psi, TRUE, fFound, (LPARAM)&sisi);
            if (!fFound)
            {
                psia->InsertAt(iInsert, psi);
            }
        }
    }
    catch (CAtlException &e)
    {
        hr = e.m_hr;
    }
    return hr;
}

HRESULT CTypeCategorizer::_CreateGroupForTypeArray(PCWSTR pszType, const CSolutionItemArray &sia, REFIID riid, __deref_out void **ppv)
{
    SHFILEINFO shfi = {0};
    HRESULT hr = SHGetFileInfo(pszType, FILE_ATTRIBUTE_NORMAL, &shfi, sizeof(shfi), SHGFI_TYPENAME | SHGFI_USEFILEATTRIBUTES) ? S_OK : E_FAIL;
    if (SUCCEEDED(hr))
    {
        CComPtr<IUnknownArray> spua;
        hr = CUnknownArray::s_CreateInstance((const CAtlArray<CComPtr<IUnknown>> &)sia, IID_PPV_ARGS(&spua));
        if (SUCCEEDED(hr))
        {
            CComPtr<ISolutionItemGroup> spsig;
            hr = CreateSolutionItemGroup(shfi.szTypeName, spua, riid, ppv);
        }
    }
    return hr;
}

STDMETHODIMP CTypeCategorizer::GetCategories(REFIID riid, __deref_out void **ppv)
{
    HRESULT hr = S_OK;
    CAtlArray<CComPtr<ISolutionItemGroup>> rgGroups;
    for (POSITION pos = _mapTypes.GetStartPosition(); pos && SUCCEEDED(hr);)
    {
        CAtlMap<CString, CSolutionItemArray>::CPair *pPair = _mapTypes.GetNext(pos);
        CComPtr<ISolutionItemGroup> spsig;
        hr = _CreateGroupForTypeArray(pPair->m_key, pPair->m_value, IID_PPV_ARGS(&spsig));
        if (SUCCEEDED(hr))
        {
            try
            {
                rgGroups.Add(spsig);
            }
            catch (CAtlException &e)
            {
                hr = e.m_hr;
            }
        }
    }

    if (SUCCEEDED(hr))
    {
        hr = CUnknownArray::s_CreateInstance((const CAtlArray<CComPtr<IUnknown>> &)rgGroups, riid, ppv);
    }
    return hr;
}
